#!/bin/bash

set -eu
# set -x


####################################################################################
#basic function                                                                    #
####################################################################################
function print_usage() {
    echo "Usage:"
    echo "    speed_profiling.sh [-m OM_MODEL] [-o OUTPUT_DIR] [-i INPUT_FILES] [-p INPUT_SHAPE] [-s OUTPUT_SIZE] [-r ROUND] [-d DEVICE_ID] [-n NPU_TYPE] [--profiler]"
    echo "Description:"
    echo "    Required parameters: "
    echo "        OM_MODEL, the path of input model"
    echo "        OUTPUT_DIR, the output dir for profling results"
    echo "    Optional parameters: "
    echo "        INPUT_FILES, the input files for model, only required for dynamic models"
    echo "        INPUT_SHAPE, the input shape for model, only required for dynamic models"
    echo "        OUTPUT_SIZE, the output size for model, only required for dynamic models"
    echo "        ROUND, round number for profiling"
    echo "        DEVICE_ID, device id for npu"
    echo "        NPU_TYPE, type for npu, etc: 310,710"
    echo "Sample:"
    echo "    Static model:"
    echo "        speed_profiling.sh -m ./debug.om -o ./profiling"
    echo "        speed_profiling.sh -m ./debug.om -o ./profiling -n 710"
    echo "    Dynamic model:"
    echo "        speed_profiling.sh -m ./debug.om -o ./profiling -i ./input_files/input.bin -p image:1,3,160,160 -s 1000000"
    echo "    Dynamic model host profiling:"
    echo "        speed_profiling.sh -m ./debug.om -o ./profiling -i ./input_files/input.bin -p image:1,3,160,160 -s 1000000 --profiler"
    exit 0
}

# get opt
OM_MODEL=''
OUTPUT_DIR=''
INPUT_FILES=''
INPUT_SHAPE=''
OUTPUT_SIZE=''
DEVICE_ID=0
ROUND=10
NPU_TYPE=310
OPEN_PROFILER=''

while getopts 'm:o:i:r:d:p:s:n:-:h' OPT; do
    # support long options
    if [ "$OPT" = "-" ]; then
        OPT="${OPTARG%%=*}"      # extract log option name
        OPTARG="${OPTARG#$OPT}"  # extract long option argument
        OPTARG="${OPTARG#=}"     # if long option argument, remove assigning `=`
    fi
    case $OPT in
        m) OM_MODEL="$OPTARG";;
        o) OUTPUT_DIR="$OPTARG";;
        i) INPUT_FILES="$OPTARG";;
        r) ROUND="$OPTARG";;
        d) DEVICE_ID="$OPTARG";;
        p) INPUT_SHAPE="$OPTARG";;
        s) OUTPUT_SIZE="$OPTARG";;
        n) NPU_TYPE="$OPTARG";;
        profiler) OPEN_PROFILER=1;;
        h) print_usage;;
        ?) print_usage;;
        *) print_usage;;
    esac
done


if [ ! $OM_MODEL ] || [ ! $OUTPUT_DIR ];then
    print_usage
fi


# init env
OM_MODEL=`realpath $OM_MODEL`
OUTPUT_DIR=`realpath $OUTPUT_DIR`
mkdir -p $OUTPUT_DIR
rm -rf $OUTPUT_DIR/*
ROOT_DIR=`pwd`
ARCH=`arch`
RUN_FILE=${OUTPUT_DIR}/run
ENV_FILE=$ROOT_DIR/speed_profiling/env.sh
source $ENV_FILE
INSTALL_PATH=`echo $install_path`

# adapt for new version CANN
if [ -d ${INSTALL_PATH}/x86_64-linux/toolkit/tools/profiler/bin ];
then
    BIN_DIR=${INSTALL_PATH}/x86_64-linux/toolkit/tools/profiler/bin
else
    BIN_DIR=${INSTALL_PATH}/toolkit/tools/profiler/bin
fi
if [ -d ${INSTALL_PATH}/x86_64-linux/toolkit/tools/profiler/profiler_tool/analysis/msprof/ ];
then
    ANALY_DIR=${INSTALL_PATH}/x86_64-linux/toolkit/tools/profiler/profiler_tool/analysis/msprof/
else
    ANALY_DIR=${INSTALL_PATH}/toolkit/tools/profiler/profiler_tool/analysis/msprof/
fi

BIN_PATH=${ROOT_DIR}/bins/msame.${ARCH}
ADA_SCRIPT=${ROOT_DIR}/speed_profiling/host_profiling.py
chmod a+x $BIN_PATH

####################################################################################
# # profiling function                                                             #
####################################################################################
function build_env() {
    echo "install_path=$INSTALL_PATH" > $RUN_FILE
    echo 'export PATH=/usr/local/python3.7.5/bin:${install_path}/atc/ccec_compiler/bin:${install_path}/atc/bin:$PATH
if [ ! `echo $PYTHONPATH` ];then
    export PYTHONPATH=${install_path}/atc/python/site-packages
else
    export PYTHONPATH=${install_path}/atc/python/site-packages:$PYTHONPATH
fi
export LD_LIBRARY_PATH=${install_path}/atc/lib64:${install_path}/acllib/lib64:$LD_LIBRARY_PATH
export ASCEND_OPP_PATH=${install_path}/opp' >> $RUN_FILE
}


function build_run_static() {
    echo "$BIN_PATH --loop $ROUND --model $OM_MODEL --device $DEVICE_ID --output $OUTPUT_DIR" >> $RUN_FILE
}


function build_run_dynamic() {
    echo "$BIN_PATH --loop $ROUND --model $OM_MODEL --device $DEVICE_ID --output $OUTPUT_DIR --input \"$INPUT_FILES\" --dymShape \"$INPUT_SHAPE\" --outputSize \"$OUTPUT_SIZE\"" >> $RUN_FILE
}

function build_run_dynamic_host() {
    echo "export GE_PROFILING_TO_STD_OUT=1" >> $RUN_FILE
    echo "$BIN_PATH --model $OM_MODEL --device $DEVICE_ID --output $OUTPUT_DIR --input \"$INPUT_FILES\" --dymShape \"$INPUT_SHAPE\" --outputSize \"$OUTPUT_SIZE\" --profiler true" >> $RUN_FILE
}

function get_profiling_name() {
    for f in `ls ${OUTPUT_DIR}`;
    do
        echo $f
        if [[ $f =~ 'JOB' ]];then
            outname=${OUTPUT_DIR}/$f
            return
        elif [[ $f =~ 'PROF' ]];then
            outname=${OUTPUT_DIR}/$f
            return
        fi
    done
}


function profiling() {
    chmod 777 $RUN_FILE

    # host dynamic model profiling
    if [ $OPEN_PROFILER ]; then
        cd $OUTPUT_DIR
        ${RUN_FILE} &> profiler.log
        python3.7 $ADA_SCRIPT profiler.log
        return
    fi

    # device model profiling
    cd $BIN_DIR
    if [ $NPU_TYPE == '310' ];then
        ./msprof --output=$OUTPUT_DIR --application=$RUN_FILE --sys-hardware-mem=on --sys-cpu-profiling=on --sys-profiling=on --sys-pid-profiling=on --sys-io-profiling=on --dvpp-profiling=on
    elif [ $NPU_TYPE == '710' ];then
        ./msprof --output=$OUTPUT_DIR --application=$RUN_FILE --sys-hardware-mem=on --sys-cpu-profiling=on --sys-profiling=on --sys-pid-profiling=on --dvpp-profiling=on
    else
        echo 'Error: only support for NPU 310/710'
        return
    fi

    cd $ANALY_DIR
    outname='fdsfasdf'
    get_profiling_name
    python3.7 msprof.py import -dir $outname
    python3.7 msprof.py export summary -dir $outname
    python3.7 msprof.py export timeline -dir $outname --iteration-id 1
    # 在chrome的地址栏输入chrome://tracing/加载打点数据查看打点图
}


function clean_dir() {
    rm $RUN_FILE
    rm -rf $OUTPUT_DIR/20*
    # rm $OUTPUT_DIR/*.*
}


####################################################################################
# # main process                                                                   #
####################################################################################
build_env

if [ ! $INPUT_FILES ] || [ ! $INPUT_SHAPE ] || [ ! $OUTPUT_SIZE ]; then
    build_run_static
else
    if [ $OPEN_PROFILER ]; then
        build_run_dynamic_host
    else
        build_run_dynamic
    fi
fi

profiling
clean_dir
